// bookreader.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "bookreader.h"

#include "zlib\zlib.h"
#include "bzip2\bzlib.h"

static inline void be2le(unsigned long& iVal)
{
	iVal = ((iVal & 0x000000ff) << 24) + 
		   ((iVal & 0x0000ff00) << 8) + 
		   ((iVal & 0x00ff0000) >> 8) + 
		   ((iVal & 0xff000000) >> 24);
}

static inline long ceildiv(long _Numerator, long _Denominator)
{
	return (int)ceil(double(_Numerator) / _Denominator);
}


//////////////////////////////////////////////////////////////////////////
// CBookReader

CBookReader::CBookReader() : m_hFile(INVALID_HANDLE_VALUE)
{
}

CBookReader::~CBookReader()
{
	if (m_hFile != INVALID_HANDLE_VALUE)
		CloseHandle(m_hFile);
}

bool CBookReader::Load(std::string strFileName)
{
	if (m_hFile != INVALID_HANDLE_VALUE)
	{
		CloseHandle(m_hFile);
		m_hFile = INVALID_HANDLE_VALUE;
	}

	m_vTableRec.clear();
	m_vChapterName.clear();
	m_vTailBlock.clear();
	m_vTailRec.clear();

	CHandle hFile(CreateFileA(strFileName.c_str(), 
							  GENERIC_READ, 
							  FILE_SHARE_READ, 
							  NULL, 
							  OPEN_EXISTING, 
							  FILE_ATTRIBUTE_NORMAL, 
							  NULL));
	if (hFile == INVALID_HANDLE_VALUE)
		return false;

	// read file header
	DWORD dwTemp = 0;
	if (!ReadFile(hFile, &m_SnbHeader, sizeof(SNBHEADER), &dwTemp, NULL))
	{
		return false;
	}

	be2le(m_SnbHeader.dwBlockSize);
	be2le(m_SnbHeader.dwBackColor);
	be2le(m_SnbHeader.dwTextColor);
	be2le(m_SnbHeader.dwChapterCount);
	be2le(m_SnbHeader.dwTableSize);
	be2le(m_SnbHeader.dwTableCompressed);
	be2le(m_SnbHeader.dwImagesSize);
	be2le(m_SnbHeader.dwContentSize);
	be2le(m_SnbHeader.dwReserved);

	if (m_SnbHeader.dwSignature != MAKEFOURCC('S', 'N', 'B', 'P') ||
		m_SnbHeader.dwStrucVersion != MAKEFOURCC('0', '0', '0', 'B'))

	{
		return false;
	}


	// read table record
	int nTitleLen = m_SnbHeader.dwTableCompressed;
	boost::scoped_array<char> pTitleBuf(new char[nTitleLen]);

	if (!ReadFile(hFile, (void*)pTitleBuf.get(), nTitleLen, &dwTemp, NULL))
	{
		return false;
	}

	LONG lDecode = m_SnbHeader.dwTableSize;
	boost::scoped_array<char> pTitleDecode(new char[m_SnbHeader.dwTableSize]);
	int ret = uncompress((Bytef *)pTitleDecode.get(), 
						 (uLongf *)&lDecode, 
						 (Bytef *)pTitleBuf.get(),
						 (uLongf)nTitleLen);
	if (ret != 0)
	{
		return false;
	}

	char* pTitlePtr = pTitleDecode.get();
	for (unsigned long i = 0; i < m_SnbHeader.dwChapterCount; ++i)
	{
		TABLEREC *pRec = new TABLEREC;
		memcpy(pRec, pTitlePtr, sizeof(TABLEREC));

		be2le(pRec->dwType);
		be2le(pRec->dwChapterNameOffset);
		be2le(pRec->dwChapterSize);

		m_vTableRec.push_back(pRec);

		pTitlePtr += sizeof(TABLEREC);
	}

	for (unsigned long i = 0; i < m_SnbHeader.dwChapterCount; ++i)
	{
		char *p = pTitlePtr + m_vTableRec[i].dwChapterNameOffset;

		m_vChapterName.push_back(new std::string(p));
	}

	// reader tailer
	DWORD dwFileSize = GetFileSize(hFile, &dwTemp);		
	LONG lTemp = 0;
	ret = SetFilePointer(hFile, dwFileSize - sizeof(SNBTAILER), &lTemp, FILE_BEGIN);
	if (ret == INVALID_SET_FILE_POINTER)
	{
		return false;
	}

	if (!ReadFile(hFile, &m_SnbTailer, sizeof(SNBTAILER), &dwTemp, NULL))
	{
		return false;
	}


	be2le(m_SnbTailer.dwBlockIdxOffset);
	be2le(m_SnbTailer.dwBlockIdxSize);

	if (m_SnbTailer.dwSignature != MAKEFOURCC('S', 'N', 'B', 'P') ||
		m_SnbTailer.dwStrucVersion != MAKEFOURCC('0', '0', '0', 'B'))
	{
		return false;
	}


	lTemp = 0;
	ret = SetFilePointer(hFile, m_SnbTailer.dwBlockIdxOffset, &lTemp, FILE_BEGIN);
	if (ret == INVALID_SET_FILE_POINTER)
	{
		return false;
	}

	boost::scoped_array<char> pTailBuf(new char[m_SnbTailer.dwBlockIdxSize]);
	if (!ReadFile(hFile, pTailBuf.get(), m_SnbTailer.dwBlockIdxSize, &dwTemp, NULL))
	{
		return false;
	}


	m_nImagesBlockCount = ceildiv(m_SnbHeader.dwImagesSize, m_SnbHeader.dwBlockSize);
	m_nContentBlockCount = ceildiv(m_SnbHeader.dwContentSize, m_SnbHeader.dwBlockSize);

	int tailfullsize = (m_nImagesBlockCount + m_nContentBlockCount) * sizeof(LONG) + 
		m_SnbHeader.dwChapterCount * sizeof(BLOCKIDXREC);
	
	lDecode = tailfullsize;
	boost::scoped_array<char> pTailDecode(new char[tailfullsize]);
	ret = uncompress((Bytef *)pTailDecode.get(), 
					 (uLongf *)&lDecode,
					 (Bytef *)pTailBuf.get(),
					 (uLongf)m_SnbTailer.dwBlockIdxSize);

	if (ret != 0)
	{
		return false;
	}

	char* pTailPtr = pTailDecode.get();
	for (int i = 0; i < m_nImagesBlockCount + m_nContentBlockCount; ++i)
	{
		unsigned long block = *(LONG *)pTailPtr;
		
		be2le(block);
		m_vTailBlock.push_back(block);

		pTailPtr += sizeof(LONG);
	}

	for (unsigned long i = 0; i < m_SnbHeader.dwChapterCount; ++i)
	{
		BLOCKIDXREC *pRec = new BLOCKIDXREC;
		memcpy(pRec, pTailPtr, sizeof(BLOCKIDXREC));

		be2le(pRec->dwBlockIdx);
		be2le(pRec->dwBlockOffset);

		m_vTailRec.push_back(pRec);

		pTailPtr += sizeof(BLOCKIDXREC);
	}

	m_hFile = hFile.Detach();

	return true;
}

bool CBookReader::GetChapter(std::string strName, std::string& strText)
{
	std::string strChapterPath = strName;

	for (unsigned long i = 0; i < m_SnbHeader.dwChapterCount; ++i)
	{
		TABLEREC rec = m_vTableRec[i];

		if (!(rec.dwType & 0x40000000))
			continue;

		std::string strFile = m_vChapterName[i];
		if (_stricmp(strFile.c_str(), strChapterPath.c_str()) == 0)
		{
			int nSize = rec.dwChapterSize;
			boost::shared_array<char> pData;

			if (ReadTextFile(m_vTailRec[i].dwBlockIdx, m_vTailRec[i].dwBlockOffset, nSize, pData))
			{
				strText = pData.get();
				return true;
			}

			return false;
		}
	}

	return false;
}

bool CBookReader::GetImage(std::string strName, boost::shared_array<char>& pImageData, int& nSize)
{
	std::string strImagePath = strName;

	for (unsigned long i = 0; i < m_SnbHeader.dwChapterCount; ++i)
	{
		TABLEREC rec = m_vTableRec[i];

		if (rec.dwType & 0x40000000)
			continue;

		std::string strFile = m_vChapterName[i];
		if (_stricmp(strFile.c_str(), strImagePath.c_str()) == 0)
		{
			nSize = rec.dwChapterSize;
			return ReadImageFile(m_vTailRec[i].dwBlockIdx, 
				m_vTailRec[i].dwBlockOffset,
				nSize,
				pImageData);
		}
	}

	return false;
}

bool CBookReader::Extract(std::string strPath)
{
	if (m_hFile == INVALID_HANDLE_VALUE)
		return false;

	HANDLE hFile = m_hFile;

	boost::scoped_array<char> pImagesBuffer(new char[m_SnbHeader.dwImagesSize]);
	boost::scoped_array<char> pContentBuffer(new char[m_SnbHeader.dwContentSize]);

	char *p = pImagesBuffer.get();
	for (int i = 0; i < m_nImagesBlockCount; ++i) // image
	{
		int blksize;
		if (i < m_nImagesBlockCount - 1)
			blksize = m_SnbHeader.dwBlockSize;
		else
			blksize = m_SnbHeader.dwImagesSize - i * m_SnbHeader.dwBlockSize;

		LONG lTemp = 0;
		DWORD dwTemp = 0;

		SetFilePointer(hFile, m_vTailBlock[i], &lTemp, FILE_BEGIN);
		ReadFile(hFile, p, blksize, &dwTemp, NULL);
	
		p += blksize;
	}


	p = pContentBuffer.get();
	for (int i = 0; i < m_nContentBlockCount; ++i)
	{
		int blksize;
		if (i < m_nContentBlockCount - 1)
			blksize = m_SnbHeader.dwBlockSize;
		else
			blksize = m_SnbHeader.dwContentSize - i * m_SnbHeader.dwBlockSize;

		int blksize_encoded;
		if (i < m_nContentBlockCount - 1)
			blksize_encoded = m_vTailBlock[m_nImagesBlockCount + i + 1] - 
							m_vTailBlock[m_nImagesBlockCount + i];
		else
			blksize_encoded = m_SnbTailer.dwBlockIdxOffset - m_vTailBlock[m_nImagesBlockCount + i];

		LONG lTemp = 0;
		DWORD dwTemp = 0;


		boost::scoped_array<char> pTempBuf(new char[blksize_encoded]);

		SetFilePointer(hFile, m_vTailBlock[m_nImagesBlockCount + i], &lTemp, FILE_BEGIN);
		ReadFile(hFile, pTempBuf.get(), blksize_encoded, &dwTemp, NULL);

		if (blksize != blksize_encoded)
		{
			unsigned int uTemp = blksize;
			::BZ2_bzBuffToBuffDecompress(p, &uTemp, pTempBuf.get(), blksize_encoded, 0, 0);
		}
		else
		{
			memcpy(p, pTempBuf.get(), blksize);
		}

		p += blksize;
	}


	std::string strStorage = strPath;
	std::string str = strStorage;
	
	str += "snbc";
	CreateDirectoryA(str.c_str(), NULL);
	
	str += "\\images";
	CreateDirectoryA(str.c_str(), NULL);
	
	str = strStorage;
	str += "snbf";
	CreateDirectoryA(str.c_str(), NULL);


	for (unsigned long i = 0; i < m_SnbHeader.dwChapterCount; ++i)
	{
		TABLEREC rec = m_vTableRec[i];
		bool bIsImage = !(rec.dwType & 0x40000000);

		std::string strFile = strStorage;
		strFile += m_vChapterName[i];

		std::replace(strFile.begin(), strFile.end(), '/', '\\');
		CHandle hOut(CreateFileA(strFile.c_str(), 
								 GENERIC_ALL, 
								 0, 
								 NULL, 
								 CREATE_ALWAYS, 
								 FILE_ATTRIBUTE_NORMAL, 
								 NULL));

		DWORD dwTemp = 0;

		if (bIsImage)
		{
			char *p = pImagesBuffer.get() + m_vTailRec[i].dwBlockIdx * m_SnbHeader.dwBlockSize + 
					m_vTailRec[i].dwBlockOffset;

			WriteFile(hOut, p, rec.dwChapterSize, &dwTemp, NULL);
		}
		else
		{
			char *p = pContentBuffer.get() + (m_vTailRec[i].dwBlockIdx - m_nImagesBlockCount) * m_SnbHeader.dwBlockSize + 
					m_vTailRec[i].dwBlockOffset;

			WriteFile(hOut, p, rec.dwChapterSize, &dwTemp, NULL);
		}
	}

	return true;
}

bool CBookReader::ReadImageFile(int nBlockIndex, int nBlockOffset, int nLength, boost::shared_array<char>& pOutData)
{
	HANDLE hFile = m_hFile;
	if (hFile == INVALID_HANDLE_VALUE)
		return NULL;

	if (nBlockIndex < 0 || nLength <= 0)
		return NULL;

	DWORD nEnd = nBlockIndex * m_SnbHeader.dwTableSize + (nBlockOffset + nLength);
	if (nEnd > m_SnbHeader.dwImagesSize)
		return false;


	boost::shared_array<char> pData(new char[nLength]);
	int nSeekPos = m_vTailBlock[nBlockIndex] + nBlockOffset;
	
	int nRead = m_SnbHeader.dwBlockSize - nBlockOffset;
	char *p = pData.get();

	int nRemain = nLength;
	while (nRemain > 0)
	{
		if (nRemain < nRead)
			nRead = nRemain;

		LONG lTemp = 0;
		DWORD dwTemp = 0;

		if (INVALID_SET_FILE_POINTER ==	SetFilePointer(hFile, nSeekPos, &lTemp, FILE_BEGIN))
		{
			return false;
		}

		if (!ReadFile(hFile, p, nRead, &dwTemp, NULL))
		{
			return false;
		}

		p += nRead;
		nRemain -= nRead;

		nSeekPos = m_vTailBlock[++nBlockIndex];
		nRead = m_SnbHeader.dwBlockSize;
	}

	pOutData = pData;

	return true;
}

bool CBookReader::ReadTextFile(int nBlockIndex, int nBlockOffset, int nLength, boost::shared_array<char>& pOutData)
{
	HANDLE hFile = m_hFile;
	if (hFile == INVALID_HANDLE_VALUE)
		return false;

	if (nBlockIndex < 0 || nLength <= 0)
		return false;

	DWORD nEnd = nBlockIndex * m_SnbHeader.dwTableSize + (nBlockOffset + nLength);
	if (nEnd > m_SnbHeader.dwContentSize)
		return false;


	boost::shared_array<char> pData(new char[nLength + 1]);
	pData[nLength] = 0;

	int nRead = m_SnbHeader.dwBlockSize - nBlockOffset;
	char *p = pData.get();


	boost::scoped_array<char> pBlankBuf(new char[m_SnbHeader.dwBlockSize]);
	boost::pool<> pool(sizeof(char));

	int nRemain = nLength;
	while (nRemain > 0)
	{
		int nSeekPos = m_vTailBlock[nBlockIndex];

		if (nRemain < nRead)
			nRead = nRemain;

		int blksize;
		if (nBlockIndex < m_nContentBlockCount - 1)
			blksize = m_SnbHeader.dwBlockSize;
		else
			blksize = m_SnbHeader.dwContentSize - 
					(nBlockIndex - m_nImagesBlockCount) * m_SnbHeader.dwBlockSize;

		int encblksize;
		if ((nBlockIndex - m_nImagesBlockCount) < m_nContentBlockCount - 1)
			encblksize = m_vTailBlock[nBlockIndex + 1] - m_vTailBlock[nBlockIndex];
		else
			encblksize = m_SnbTailer.dwBlockIdxOffset - m_vTailBlock[nBlockIndex];

		LONG lTemp = 0;
		DWORD dwTemp = 0;


		char *pTempBuf = (char *)pool.ordered_malloc(encblksize);
		BOOST_ASSERT(pTempBuf);

		SetFilePointer(hFile, nSeekPos, &lTemp, FILE_BEGIN);
		ReadFile(hFile, pTempBuf, encblksize, &dwTemp, NULL);

		if (encblksize != nRead) //fuuu
		{
			unsigned int uTemp = blksize;
			::BZ2_bzBuffToBuffDecompress(pBlankBuf.get(), &uTemp, pTempBuf, encblksize, 0, 0);

			memcpy(p, pBlankBuf.get() + nBlockOffset, nRead);
		}
		else
		{
			memcpy(p, pTempBuf, nRead);
		}

		pool.free(pTempBuf);

		p += nRead;
		nRemain -= nRead;

		nRead = m_SnbHeader.dwBlockSize;
		nBlockOffset = 0;

		++nBlockIndex;
	}

	pOutData = pData;

	return true;
}

#if 0

int _tmain(int argc, _TCHAR* argv[])
{
	CSnbReader reader;
	if (!reader.Load("c:\\test2.snb"))
	{
		std::cout << "failed open" << std::endl;
		exit(0);
	}

	reader.Extract("c:\\bambook2\\");

	return 0;
}

#endif